
import time

from threading import Lock

import rpyc
from rpyc.utils import server

import dexen
from dexen.system import constants

#===============================================================================
# Globals
#===============================================================================
CONNECTED = -1
DISCONNECTED = -2
ALLOCATED = -3

slave_prefix = "slave"
slaves = []
lock = Lock() 

#===============================================================================
# The helper and decorator methods
#===============================================================================
def synchronized(f):
    """A decorator to synchronize all accesses from remote clients.
    
    This decorator function is specifically used to protect read/write acceses 
    to the globally shared'slaves' data. Any service method accesses 'slaves'
    should be decorated.
    
    """
    def new(*args, **kwargs):
        try:
            lock.acquire()
            res = f(*args, **kwargs)
            return res
        finally:
            lock.release()
    return new

def localized(f):
    def new(*args, **kwargs):
        _args = list(args)
        print _args
        for i, arg in enumerate(_args):
            if i != 0:
                _args[i] = rpyc.classic.obtain(arg)
            
        for key in kwargs.keys():
            kwargs[key] = rpyc.classic.obtain(kwargs[key])
            
        return f(*_args, **kwargs)
    return new

def generate_name(id):
    return slave_prefix + str(id)

def prepare_slave_info(id):
    output = (generate_name(id),)    
    return output

#===============================================================================
# The Service Class
#===============================================================================
class NamingService(dexen.Service):
    """A service that names slaves and manages their allocation."""
    
    ALIASES = [constants.NS_NAME]
        
    @synchronized    
    def on_disconnect(self):
        """Unregister the slave being disconnected."""
        slaves[self._id] = DISCONNECTED
        super(NamingService, self).on_disconnect(self)
    
    def _connect_slave(self):
        """Insert the slave in to the 'slaves' list.
        
        Returns the assigned id.
        
        """
        for i, status in enumerate(slaves):
            if status == DISCONNECTED:
                slaves[i] = CONNECTED
                return i
        slaves.append(CONNECTED)
        return len(slaves)-1
    
    @synchronized
    def register_slave(self):
        """Register the slave by giving it a name.
        
        Notes:
        The slave names are in the form of 'slave[0-9]*'.
        
        Returns a tuple (assigned_name, ) to the slave. 
        
        """
        self._id = self._connect_slave()
        self._slave_name = generate_name(self._id) 
        return prepare_slave_info(self._id)
    
    @synchronized
    def get_free_slave_names(self):
        """Allocate slaves to the Job Dispatcher module.
        
        Reserves all the slaves that are available.
        
        Returns the names of the available slaves.
        
        """
        res = []
        for i, status in enumerate(slaves):
            if status == CONNECTED:
                slaves[i] = ALLOCATED
                res.append(generate_name(i))
        return res
 
        
def main():
    svr = server.ThreadedServer(service=NamingService, auto_register=True,
                                protocol_config={"allow_all_attrs":True})
    svr.start()
    
if __name__ == '__main__':
    main()

